package cn.echoparrot.domain.entity;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.util.Assert;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 实现长期记忆，用于对话的保存
 * 设计的基本假设：
 * 1. 对聊天记录的成本因素考虑权重远远大于可靠性的考虑权重；
 * 2. 尽量减少内存的压力与IO的压力；默认参数不保证宕机情况下消息可靠存储，如果需要可靠存储也可将batchOutMessages参数设置为1。
 * 3. 满足上述成本要求的前提下，尽量保证聊天的质量效果。
 * @see org.springframework.ai.chat.memory.MessageWindowChatMemory
 * @see # org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
 *
 * ChatMemory里不记录SystemMessage，SystemMessage的权重为0，不参与长期记忆的计算。
 * SystemMessage 的模板，是系统单独记录的，而里面的变量，如果有需要，可以记录在 Conversation里。例如中期记忆（消息摘要）
 * 通过自定义一个Advisor向Prompt注入SystemMessage，参考PromptChatMemoryAdvisor。
 * @author Able
 */
public class SummaryChatMemory implements ChatMemory {

    private final Map<String, SummaryConversation> conversations = new ConcurrentHashMap<>();

    private final ChatMemoryRepository chatMemoryRepository;

    public SummaryChatMemory(ChatMemoryRepository chatMemoryRepository) {
        this.chatMemoryRepository = chatMemoryRepository;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        Assert.notNull(messages, "messages cannot be null");
        Assert.noNullElements(messages, "messages cannot contain null elements");
        SummaryConversation conversation = conversations.computeIfAbsent(conversationId,this::provideConversation) ;

        conversation.add(messages);
    }

    @Override
    public List<Message> get(String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        SummaryConversation conversation =conversations.computeIfAbsent(conversationId,this::provideConversation) ;
        return conversation.get();
    }

    @Override
    public void clear(String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        SummaryConversation conversation =conversations.remove(conversationId);
        conversation.flush();
    }

    private SummaryConversation provideConversation(String conversationId){
        // 从数据库获取，构建Conversation对象
        List<Message> stored =chatMemoryRepository.findByConversationId(conversationId);
        var conversation = new SummaryConversation(conversationId,stored, chatMemoryRepository, MemorySettings.DEFAULT);
        return conversation;
    }
}
